package top.yudoge.dogeeditor.llm.agent

import cn.hutool.core.util.StrUtil
import dev.langchain4j.data.message.UserMessage
import dev.langchain4j.memory.chat.MessageWindowChatMemory
import dev.langchain4j.service.AiServices
import dev.langchain4j.service.MemoryId
import dev.langchain4j.service.SystemMessage
import dev.langchain4j.service.TokenStream
import org.springframework.beans.factory.annotation.Autowired
import org.springframework.stereotype.Component
import top.yudoge.dogeeditor.commons.cache.Cache
import top.yudoge.dogeeditor.db.api.DataSourceMetaDataApi
import top.yudoge.dogeeditor.llm.constants.SessionMemoryKeyConstants.Companion.DB_NAME
import top.yudoge.dogeeditor.llm.constants.SessionMemoryKeyConstants.Companion.DS_NAME
import top.yudoge.dogeeditor.llm.constants.SessionMemoryKeyConstants.Companion.SCHEMA_NAME
import top.yudoge.dogeeditor.llm.exceptions.RewriteException
import top.yudoge.dogeeditor.llm.memory.SessionMemory
import top.yudoge.dogeeditor.llm.config.models.ModelRegistry
import top.yudoge.dogeeditor.llm.objects.query.RewriteQuery
import top.yudoge.dogeeditor.llm.util.tokenizer.Tokenizer
import top.yudoge.dogeeditor.llm.tools.DBTools
import java.util.*
import javax.annotation.PostConstruct

//@Deprecated("Not used anymore")
//@Component
//class LangChain4JRewriteAgent : RewriteAgent {
//
//    @Autowired
//    private lateinit var sessionMemory: SessionMemory
//
//    @Autowired
//    private val modelRegistry = ModelRegistry()
//
//    @Autowired
//    private lateinit var metaDataApi: DataSourceMetaDataApi
//
//    @Autowired
//    private lateinit var cache: Cache<String>
//
//    @Autowired
//    private lateinit var tokenizer: Tokenizer
//
//    @Autowired
//    private lateinit var tableEmbeddingService: TableEmbeddingService
//
//    private val modelNameAIServiceMap = mutableMapOf<String, RewriteAIService>()
//
//    @PostConstruct
//    fun init() {
//        modelRegistry.allStreamingChatModelNames().forEach { modelName ->
//            var streamingChatModel = modelRegistry.getStreamingChatModel(modelName)
//            val aiService = AiServices.builder(RewriteAIService::class.java)
//                .chatMemoryProvider { _ -> MessageWindowChatMemory.withMaxMessages(10) }
//                .streamingChatLanguageModel(streamingChatModel)
//                .tools(DBTools(sessionMemory, metaDataApi, cache, tokenizer, tableEmbeddingService))
//                .build()
//            modelNameAIServiceMap[modelName] = aiService
//        }
//    }
//
//    override fun rewrite(query: RewriteQuery, callback: RewriteCallback) {
//        // 0. check the query is valid
//        if (!StrUtil.isAllNotBlank(query.sessionId, query.userTask, query.dataSourceIdentifier, query.databaseName, query.schemaName)) {
//            callback.onError(RewriteException("query is not valid"))
//            return
//        }
//
//        if (StrUtil.isBlank(query.modelName) || !modelNameAIServiceMap.containsKey(query.modelName)) {
//            callback.onError(RewriteException("model ${query.modelName} is not found"))
//            return
//        }
//
//        // 1. check the session is created
//        if (!sessionMemory.containsSession(query.sessionId)) {
//            callback.onError(RewriteException("session ${query.sessionId} is not found"))
//            return
//        }
//
//        // 2. put datasource information to session memory
//        try {
//            sessionMemory.putMultiple(query.sessionId,
//                DS_NAME to query.dataSourceIdentifier,
//                DB_NAME to query.databaseName,
//                SCHEMA_NAME to query.schemaName
//            )
//
//            // 3. call ai service
//            callAIService(query.sessionId, query, callback)
//
//        } finally {
//            // 4. release for safety purpose
//            sessionMemory.removeMultiple(query.sessionId, DS_NAME, DB_NAME, SCHEMA_NAME)
//        }
//
//    }
//
//    private fun callAIService(memoryId: String, query: RewriteQuery, callback: RewriteCallback) {
//
//        var tokenStream = modelNameAIServiceMap[query.modelName]!!.chat(memoryId, buildMessages(query))
//
//        tokenStream.onError {
//            callback.onError(it)
//        }
//
//        tokenStream.onCompleteResponse {
//            callback.onComplete()
//        }
//
//        tokenStream.onPartialResponse {
//            print(it)
//            callback.onTokenGenerated(it)
//        }
//
//        tokenStream.onToolExecuted { it ->
//            println(it)
//        }
//        tokenStream.start()
//    }
//
//    private fun buildMessages(request: RewriteQuery): List<UserMessage> {
//        return Collections.unmodifiableList(mutableListOf<UserMessage>().apply {
//            val tableSelected = request.tableSelected
//            if (tableSelected.isNotEmpty()) {
//                add(UserMessage("Selected table:\n ${tableSelected.joinToString(",")}\n"))
//            }
//            if (StrUtil.isNotBlank(request.code)) {
//                add(UserMessage("User code:\n${request.code}\n"))
//            }
//
//            add(UserMessage("User task description:\n${request.userTask}\n"))
//        })
//    }

//}


internal interface RewriteAIService {
    @SystemMessage(
        "You are an excellent SQL engineer. Your task is helping users rewrite their SQL script." +
                "Your output should include your thought process and only one SQL script as the answer." +
                "Please use the same language in User Task"
    )
    fun chat(
        @MemoryId memoryId: String,
        @dev.langchain4j.service.UserMessage userMessages: List<UserMessage>
    ): TokenStream
}
